package lib

import (
	"algorithms/ds"
)

func init() {
	Probs = append(Probs, Problem{
		Num:         250,
		Discription: "统计同值子树",
		Level:       1,
		Labels: map[string]int{
			"二叉树": 1,
			"递归":  1,
			"回溯":  1,
		},
	})
}

func CountUnivalSubtrees(root *ds.TreeNode) int {
	if root == nil {
		return 0
	}

	count := 0
	var isUnival func(*ds.TreeNode) bool
	isUnival = func(node *ds.TreeNode) bool {
		if node == nil {
			return false
		}

		if node.Left == nil && node.Right == nil {
			count++
			return true
		}

		isLeftUnival := isUnival(node.Left)
		isRightUnival := isUnival(node.Right)
		if isLeftUnival && node.Right == nil && node.Left.Val == node.Val {
			count++
			return true
		}
		if isRightUnival && node.Left == nil && node.Right.Val == node.Val {
			count++
			return true
		}
		if isLeftUnival && isRightUnival && node.Left.Val == node.Val && node.Right.Val == node.Val {
			count++
			return true
		}

		return false
	}

	isUnival(root)
	return count
}
